from datasets import load_dataset
from transformers import AutoTokenizer
from sparseml.transformers import SparseAutoModelForCausalLM, oneshot
import torch

ALPACA_TEMPLATE = {
    "prompt_input": "Below is an instruction that describes a task, paired with an "
    "input that provides further context. Write a response that appropriately "
    "completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n"
    "{input}\n\n### Response:\n",
    "prompt_no_input": "Below is an instruction that describes a task. Write a "
    "response that appropriately completes the request.\n\n### Instruction:\n{"
    "instruction}\n\n### Response:\n",
}

recipe_stub = """
test_stage:
  obcq_modifiers:
    SparseGPTModifier:
      sparsity: 0.5
      block_size: 128
      sequential_update: False
      quantize: False
      targets: [
        "re:model.layers.\\\d+$"
      ]
"""

MODEL_STUB = "TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T"
OUTPUT_DIR = "./test_oneshot_tokenized_input"
NUM_CALIBRATION_SAMPLES = 1024

model = SparseAutoModelForCausalLM.from_pretrained(MODEL_STUB, device_map="auto", torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(MODEL_STUB)

dataset = load_dataset("garage-bAInd/Open-Platypus")["train"]
dataset = dataset.shuffle(seed=42).select(range(256))

def preprocess_for_oneshot(sample):
    if "input" in sample:
        concat_text = ALPACA_TEMPLATE["prompt_input"].format(
            instruction=sample["instruction"], input=sample["input"]
        )
    else:
        concat_text = ALPACA_TEMPLATE["prompt_no_input"].format(
            instruction=sample["instruction"]
        )
    if "output" in sample:
        concat_text += sample["output"]

    return tokenizer(
      concat_text,
      padding=False,
      max_length=512,
      truncation=True
    )

tokenized_dataset = dataset.map(preprocess_for_oneshot, remove_columns=["input", "output", "instruction", "data_source"])

oneshot(
  model=model,
  dataset=tokenized_dataset,
  recipe=recipe_stub,
  num_calibration_samples=NUM_CALIBRATION_SAMPLES,
  output_dir=OUTPUT_DIR,
  overwrite_output_dir=True
)